from itertools import combinations
from datetime import datetime, date
import backtrader as bt
import akshare as ak
import pandas as pd
import numpy as np
import quantstats as qs
import pyfolio as pf
import webbrowser
from sqlalchemy import create_engine, distinct, or_, and_
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser
import warnings

warnings.filterwarnings('ignore')

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


def get_stocks():
    sql = "SELECT STK_CODE FROM ETF_INDEX_INFO"
    r = exec_sql(sql)
    df1 = pd.DataFrame(r, columns=['代码'])
    df2 = ak.fund_etf_spot_em()
    res = pd.merge(df1, df2, on=['代码'])
    res = res.loc[res['最新价'] < 10, :]
    stocks = res['代码'].to_list()
    # stocks = [tok for tok in stocks if tok != '512870']

    return stocks


def get_data(stocks, freq='day', source='db'):
    if freq == '1min':
        table_name = 'STOCK_MIN_DATA'
    elif freq == '15min':
        table_name = 'STOCK_15MIN_DATA'
    else:
        table_name = 'ETF_DAY_DATA'  # HIS_DAILY_DATA

    res = dict()
    date_data = set_date()
    st = date_data['st']
    et = date_data['et']
    st_d = date_data['st_d']
    et_d = date_data['et_d']

    if source == 'db':
        print('db')
        stocks_str = ','.join(f"'{tok}'" for tok in stocks)
        sql = f"SELECT code, datetime, [close], [open], low, high, volume FROM {table_name} WHERE code in ({stocks_str})"
        r = exec_sql(sql)
        df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume'])
        # df['datetime'] = df['datetime'].map(time_map)

        for code in stocks:
            if code == '399300':
                sig_df = ak.stock_zh_index_daily(symbol="sz399300")
                sig_df = sig_df.loc[(sig_df.date >= st_d) & (sig_df.date <= et_d), :]
                sig_df.rename(columns={'date': 'datetime'}, inplace=True)
                sig_df.insert(0, 'code', code)
                sig_df = sig_df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
                # sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            else:
                sig_df = df.loc[df['code'] == code]
            sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            sig_df.sort_values('datetime', inplace=True, ignore_index=True)
            sig_df.set_index('datetime', drop=False, inplace=True)
            if len(sig_df) <= 220:
                continue
            print(f'code={code}->len={len(sig_df)}')
            res[code] = sig_df
    else:
        print('internet')
        for code in stocks:
            if code == '399300':
                df = ak.stock_zh_index_daily(symbol="sz399300")
                df = df.loc[(df.date >= st_d) & (df.date <= et_d), :]
                df.rename(columns={'date': 'datetime'}, inplace=True)
            else:
                df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=st, end_date=et, adjust="qfq")
                df.rename(columns={'日期': 'datetime', '开盘': 'open', '收盘': 'close', '最高': 'high',
                                   '最低': 'low', '成交量': 'volume'}, inplace=True)

            df.insert(0, 'code', code)
            df = df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
            df['datetime'] = pd.to_datetime(df['datetime'])
            df.set_index('datetime', inplace=True)
            print(f'code={code}->len={len(df)}')
            res[code] = df

    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def time_map(x):
    t = datetime.strptime(str(x), '%Y-%m-%d')
    # t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


def set_date():
    st = "20140101"
    et = "20230816"
    st_dt = datetime.strptime(st, '%Y%m%d')
    et_dt = datetime.strptime(et, '%Y%m%d')
    st_d = st_dt.date()
    et_d = et_dt.date()
    res = {'st': st, 'et': et, 'st_d': st_d, 'et_d': et_d, 'st_dt': st_dt, 'et_dt': et_dt}

    return res


class PairsTrading(bt.Strategy):
    params = (
        ('money', 1e5),
        ('formation_period', 120),
        ('trading_period', 60),
        ('num_pairs', 100),
        ('std_dev_threshold', 1),
        ('printlog', False),
    )

    def __init__(self):
        self.order = None
        self.counter = 0
        self.pairs = list()
        self.symbols = list()
        self.trail_dict = dict()
        self.daily_pnl = dict()

        for data in self.datas:
            code = data._name
            self.trail_dict[code] = {'counter': 0}

    def log(self, txt, dt=None, doprint=True):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            content = f'{dt.isoformat()}, {txt}'
            print(content)

    def prenext(self):
        self.next()

    def downcast(self, price):
        res = self.p.money / price // 100 * 100
        return res

    def next(self):
        if self.order:
            return

        self.counter += 1
        if self.counter < self.params.formation_period:
            return

        cur_dt = self.datas[0].datetime.datetime(0)
        dt_int = int(cur_dt.strftime('%Y%m%d'))

        if dt_int not in self.daily_pnl:
            self.daily_pnl[dt_int] = round(self.broker.getvalue() - init_cash, 2)

        for data in self.datas:
            code = data._name
            try:
                dt = data.datetime.datetime(0)  # 获取当前的回测时间点
            except:
                print(f'{code}-{dt}')
            if code == '399300' or dt != cur_dt:
                continue
            elif self.trail_dict[code]['counter'] < self.p.formation_period:
                self.trail_dict[code]['counter'] = self.trail_dict[code]['counter'] + 1
                continue
            self.symbols.append(data)

        if self.counter % self.params.trading_period == 0:
            self.select_pairs()

        for data1, data2 in self.pairs:
            spread = self.calculate_spread(data1, data2)
            std_dev = np.std(spread)
            position1 = self.getposition(data1).size
            position2 = self.getposition(data2).size
            if abs(spread.iloc[-1]) > self.params.std_dev_threshold * std_dev:
                if spread.iloc[-1] > 0:
                    amount = self.downcast(data1.close[0])
                    self.order = self.buy(data1, size=amount)
                    if position2 > 0:
                        self.order = self.close(data2)
                elif spread.iloc[-1] < 0:
                    amount = self.downcast(data2.close[0])
                    self.order = self.buy(data2, size=amount)
                    if position1 > 0:
                        self.order = self.close(data1)
            elif abs(spread.iloc[-1]) <= std_dev or self.counter % self.params.trading_period == 0:
                if position1 != 0:
                    self.order = self.close(data1)
                if position2 != 0:
                    self.order = self.close(data2)

    def select_pairs(self):
        distances = {}
        for pair in combinations(self.symbols, 2):
            data1 = pd.Series(pair[0].get(size=self.params.formation_period))
            data2 = pd.Series(pair[1].get(size=self.params.formation_period))
            normalized_data1 = data1 / data1.iloc[0]
            normalized_data2 = data2 / data2.iloc[0]
            distance = np.sum((normalized_data1 - normalized_data2) ** 2)
            distances[pair] = distance
        sorted_pairs = sorted(distances.items(), key=lambda x: x[1])
        self.pairs = [pair[0] for pair in sorted_pairs[:self.params.num_pairs]]

    def calculate_spread(self, data1, data2):
        data1_prices = pd.Series(data1.get(size=self.params.formation_period))
        data2_prices = pd.Series(data2.get(size=self.params.formation_period))
        data1_normalized = data1_prices / data1_prices.iloc[0]
        data2_normalized = data2_prices / data2_prices.iloc[0]
        spread = data1_normalized - data2_normalized
        return spread


init_cash = 5e6


def run_backtrade(stocks):
    cerebro = bt.Cerebro()
    cerebro.addstrategy(PairsTrading)
    data_all = get_data(stocks, source='db')
    index = stocks[0]
    data0 = data_all[index]
    datafeed0 = bt.feeds.PandasData(dataname=data0)
    cerebro.adddata(datafeed0, name=f'{index}')

    date_data = set_date()
    st_dt = date_data['st_dt']
    et_dt = date_data['et_dt']

    for i, code in enumerate(data_all):
        if code == index:
            df = data_all[code]
            df = df.loc[:, ['close']]
            df.rename(columns={'close': '399300'}, inplace=True)
            continue
        data = data_all[code]
        datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_dt, todate=et_dt)
        cerebro.adddata(datafeed, name=code)
        print(f'{code} feeds ok')

    # 初始资金 1,000,000.0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    # cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    strat = result[0]
    # st = strat.params.st
    pyfoliozer = strat.analyzers.getbyname('pyfolio')
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    # print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

    returns, positions, transactions, gross_lev = pyfoliozer.get_pf_items()
    print(st_dt.date())
    returns = returns[returns.index.date >= st_dt.date()]
    returns.index = returns.index.date
    returns.index = pd.to_datetime(returns.index)

    print(f'returns={returns}')
    html_file = 'test.html'
    qs.reports.html(returns, benchmark=df, title='grid', output=html_file, download_filename=html_file)
    # qs.reports.html(returns, title='grid', output=html_file, download_filename=html_file)
    qs.reports.basic(returns)
    f = webbrowser.open(html_file)


def main():
    index = '399300'
    # stocks = get_stocks()
    # stocks = ['510880', '159915', '513100', '518880']
    # stocks = ['159901', '510880']
    # stocks = ['159901', '510880', '512510', '512220', '159938', '159909', '513500', '513880', '515770', '515350', '159810', '517350']
    # stocks = stocks[:200]
    # stocks = ['159909', '510170', '159915', '512220', '510880', '510300']
    #
    # stocks = ['159918', '159909', '512220', '159915', '510500', '510510', '159907', '159922', '510290', '159936',
    #            '159930', '159935', '159912', '159903', '159822', '513050', '159902', '510900', '512070', '510410']
    # stocks = ['510300', '511220', '518880', '513100']
    stocks = ['510050', '510300', '510500', '159915', '588000', '159928', '512010', '159939',
              '510230', '562910', '513500', '513100', '515030', '159995', '512660', '515790', '512790']
    # print(stocks)
    stocks.insert(0, index)
    run_backtrade(stocks)
    # s = get_data(stocks, freq='day', source='not db')
    # print(s)


if __name__ == '__main__':
    main()
